import random
import numpy as np
import math
from view.endSrc.MixedGaussianDataset import MixedGaussianDataset
from view.endSrc.GaussianClusters import GaussianClusters


class NonoverlappingClustersDataGenerator:

    def __init__(self, numOfInstances, numOfFeatures, paiList, centerRange=None):

        for p in paiList:
            assert p > 0, 'every element in paiList must be positive'

        self.n = numOfInstances
        self.d = numOfFeatures
        self.k = len(paiList)

        self.meanList = None
        self.varMatList = None
        self.labels = None  # store generated labels from recently called self.genMixedDataset
        self.mgDataset = None  # MixedGaussianDataset

        # Normalize the pai value to ensure the summation of all pais equals to one
        self.paiList = [pai / sum(paiList) for pai in paiList]
        self.numOfInstancesEach = [int(pai * self.n) for pai in self.paiList]
        assert all(n > 0 for n in self.numOfInstancesEach), 'numOfInstances is too small, please try again'

        # Ensure the sum of instances of all components is equal to the total number of instances
        if sum(self.numOfInstancesEach) != self.n:
            print('fit:', self.n - sum(self.numOfInstancesEach))
            self.numOfInstancesEach[0] += self.n - sum(self.numOfInstancesEach)

        self.centerRange = centerRange
        if centerRange is None:
            self.centerRange = [1, self.k]

        self.gridFactor = math.ceil(pow(2*self.k, 1.0/self.d))
        # print(self.k / pow(self.gridFactor, self.d))

    def genSeparableComponents(self):
        '''
        This method generate the mean and covariance given certain parameter settings for one cluster,
        which are stored in self.meanList and self.varMatList.


        Examples:
        --------
        datasetGen.genSeparableComponents()
        '''

        gridList = np.array([np.linspace(self.centerRange[0], self.centerRange[1], self.gridFactor * 2) for dim_i in range(self.d)])
        gap = gridList[0, 1] - gridList[0, 0]
        cluster_max_std = gap / 6
        cluster_max_var = cluster_max_std * cluster_max_std

        self.meanList = []
        ridsSet = []
        while len(self.meanList) < self.k:

            randomIds = [random.randint(0, self.gridFactor-1) for dim_i in range(self.d)]
            if randomIds in ridsSet:
                continue
            ridsSet.append(randomIds)

            centerBox_low = []
            centerBox_high = []

            for dim_i in range(self.d):
                rid = randomIds[dim_i]
                centerBox_low.append(gridList[dim_i, 2 * rid])
                centerBox_high.append(gridList[dim_i, 2 * rid + 1])

            meanV = np.random.uniform(centerBox_low, centerBox_high)  # create random mean vector within box constraint
            self.meanList.append(meanV)

        self.varMatList = []
        for i in range(self.k):
            cov = np.zeros((self.d, self.d))
            for j in range(self.d):
                cov[j, j] = -np.random.uniform(low=-cluster_max_var, high=0)
            self.varMatList.append(cov)

    def genMixedDataset(self):
        # firstly, generate meanList, varMatList for Gaussian distribution
        if self.meanList is None:
            self.genSeparableComponents()
        print('numOfInstances: ', self.numOfInstancesEach)

        dataset = []
        self.labels = []
        for ind in range(self.k):
            mean = self.meanList[ind]
            cov = self.varMatList[ind]
            size = self.numOfInstancesEach[ind]

            data = np.random.multivariate_normal(mean, cov, size)
            dataset.append(data)
            self.labels += ([ind] * size)
        dataset = np.concatenate(dataset, axis=0)

        print('self.k: ', self.k)
        # print('self.labels: ', self.labels)

        # build output object
        gc = GaussianClusters(self.meanList, self.varMatList, self.paiList, labels=self.labels)
        self.mgDataset = MixedGaussianDataset(dataset, self.labels, gc)

        return self.mgDataset


# if __name__ == '__main__':
#
#     ps = []
#     for i in range(10):
#         ps.append(random.uniform(1, 3))
#
#     print(ps)
#     g = NonoverlappingClustersDataGenerator(numOfInstances=2000, numOfFeatures=2, paiList=ps, centerRange=[10, 100])
#
#     g.genSeparableComponents()
#     mg = g.genMixedDataset()

